#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Split large AMOS proceedings into multiple volumes for easier printing and handling.

Each volume will have:
- Its own cover page (e.g., "AMOS 2025 Proceedings - Volume 1 of 3")
- Its own table of contents (only papers in that volume)
- Approximately 500 pages of content (configurable)
- Proper page numbering and bookmarks

Usage:
  python src/split_proceedings.py --csv data/metadata_2025.csv --output-dir data/volumes \
      --max-pages 500 --title "AMOS 2025 Proceedings" \
      --subtitle "Advanced Maui Optical and Space Surveillance Technologies Conference"
"""
from __future__ import annotations

import argparse
import csv
import os
import sys
from dataclasses import dataclass
from typing import List, Tuple

from reportlab.lib.pagesizes import A4
from reportlab.pdfgen import canvas
from reportlab.lib.units import mm
from reportlab.lib import colors

from pypdf import PdfReader, PdfWriter
from io import BytesIO


@dataclass
class Entry:
    title: str
    authors: str
    pdf_path: str
    pages: int = 0
    start_page: int = 0
    volume: int = 0
    printed_start: int = 0  # Printed page (TOC=1) where paper starts
    blanks_before: int = 0  # Blank pages inserted before this paper


def read_entries(csv_path: str) -> Tuple[int, List[Entry]]:
    entries: List[Entry] = []
    year: int = 0
    with open(csv_path, newline='', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        for row in reader:
            try:
                year = int(row.get('year') or year or 0)
            except Exception:
                pass
            title = (row.get('title') or '').strip()
            authors = (row.get('authors') or '').strip()
            pdf_path = (row.get('pdf_path') or '').strip()
            if not title or not pdf_path:
                continue
            if not os.path.exists(pdf_path):
                candidate = os.path.join(os.path.dirname(csv_path), pdf_path)
                if os.path.exists(candidate):
                    pdf_path = candidate
                else:
                    continue
            entries.append(Entry(title=title, authors=authors, pdf_path=pdf_path))
    entries.sort(key=lambda e: e.title.lower())
    return year, entries


def measure_pages(entries: List[Entry]) -> None:
    for e in entries:
        try:
            reader = PdfReader(e.pdf_path)
            e.pages = len(reader.pages)
        except Exception:
            e.pages = 0


def assign_volumes(entries: List[Entry], target_volumes: int) -> int:
    """Assign each entry to a volume using First-Fit Decreasing-like balancing.
    Returns actual volume count. Does not change per-volume ordering (keeps alphabetical).
    """
    if target_volumes <= 0:
        target_volumes = 1

    vol_pages = [0] * target_volumes
    # Assign indices in descending order of paper length
    indices = sorted(range(len(entries)), key=lambda i: entries[i].pages, reverse=True)
    for idx in indices:
        # Choose the lightest volume so far
        vol = min(range(target_volumes), key=lambda v: vol_pages[v])
        entries[idx].volume = vol + 1
        vol_pages[vol] += entries[idx].pages

    actual_volumes = sum(1 for p in vol_pages if p > 0)
    return max(actual_volumes, 1)


def generate_cover(cover_path: str, main_title: str, subtitle: str, year: int, volume_num: int, total_volumes: int) -> int:
    c = canvas.Canvas(cover_path, pagesize=A4)
    width, height = A4

    c.setFillColor(colors.HexColor('#0b3d91'))
    c.rect(0, 0, width, height, fill=1, stroke=0)

    c.setFillColor(colors.white)
    c.setFont('Helvetica-Bold', 36)
    c.drawCentredString(width/2, height*0.68, main_title)

    # Volume indicator
    c.setFont('Helvetica-Bold', 24)
    c.drawCentredString(width/2, height*0.60, f"Volume {volume_num} of {total_volumes}")

    # Subtitle
    c.setFont('Helvetica', 16)
    subtitle_max_width = width - 60 * mm
    if c.stringWidth(subtitle, 'Helvetica', 16) > subtitle_max_width:
        words = subtitle.split()
        mid = len(words) // 2
        line1 = ' '.join(words[:mid])
        line2 = ' '.join(words[mid:])
        c.drawCentredString(width/2, height*0.54, line1)
        c.drawCentredString(width/2, height*0.51, line2)
    else:
        c.drawCentredString(width/2, height*0.54, subtitle)

    c.setFont('Helvetica', 14)
    c.drawCentredString(width/2, height*0.46, f"Proceedings • {year}")

    c.setFont('Helvetica', 10)
    c.drawCentredString(width/2, height*0.08, "Compiled from AMOS Technical Library metadata")

    c.showPage()
    c.save()
    return 1


def wrap_text(text: str, max_chars: int) -> List[str]:
    lines: List[str] = []
    cur = []
    for part in text.split():
        if sum(len(x) for x in cur) + len(cur) + len(part) > max_chars and cur:
            lines.append(' '.join(cur))
            cur = [part]
        else:
            cur.append(part)
    if cur:
        lines.append(' '.join(cur))
    return lines


def generate_toc(toc_path: str, entries: List[Entry], volume_num: int) -> int:
    c = canvas.Canvas(toc_path, pagesize=A4)
    width, height = A4
    left = 20 * mm
    right = width - 20 * mm
    top = height - 20 * mm
    bottom = 20 * mm

    header_style = ('Helvetica-Bold', 20)
    title_style = ('Helvetica-Bold', 11)
    author_style = ('Helvetica', 10)

    line_height = 13
    max_title_chars = 75

    y = top
    pages_used = 1

    def new_page():
        nonlocal y
        c.showPage()
        y = top

    # Header
    c.setFont(*header_style)
    c.drawString(left, y, f'Table of Contents - Volume {volume_num}')
    y -= 24

    for e in entries:
        c.setFont(*title_style)
        title_lines = wrap_text(e.title, max_title_chars)
        for i, line in enumerate(title_lines):
            if y - line_height < bottom:
                new_page()
                pages_used += 1
                c.setFont(*header_style)
                c.drawString(left, y, f'Table of Contents - Volume {volume_num} (cont.)')
                y -= 24
                c.setFont(*title_style)
            
            if i == 0:
                page_num = str(e.printed_start if e.printed_start else 1)
                page_area_width = 40
                title_max_width = right - left - page_area_width - 10
                
                while c.stringWidth(line, *title_style) > title_max_width and len(line) > 10:
                    line = line[:-4] + '...'
                
                c.drawString(left, y, line)
                title_end_x = left + c.stringWidth(line, *title_style)
                dots_start = title_end_x + 5
                dots_end = right - page_area_width
                if dots_start < dots_end - 10:
                    c.setDash(1, 3)
                    c.line(dots_start, y + 3, dots_end, y + 3)
                    c.setDash()
                c.drawRightString(right, y, page_num)
            else:
                c.drawString(left + 5, y, line)
            y -= line_height
        
        # Authors
        if e.authors:
            c.setFont(*author_style)
            authors_line = e.authors
            if ';' in authors_line:
                authors_line = authors_line.split(';')[0].strip() + ' et al.'
            elif ',' in authors_line and authors_line.count(',') > 1:
                parts = authors_line.split(',')
                if len(parts) >= 2:
                    authors_line = f"{parts[0].strip()}, {parts[1].strip()} et al."
                else:
                    authors_line = parts[0].strip() + ' et al.'
            
            max_author_width = right - left - 10
            while c.stringWidth(authors_line, *author_style) > max_author_width and len(authors_line) > 20:
                authors_line = authors_line[:len(authors_line)-4] + '...'
            
            if y - line_height < bottom:
                new_page()
                pages_used += 1
                c.setFont(*header_style)
                c.drawString(left, y, f'Table of Contents - Volume {volume_num} (cont.)')
                y -= 24
                c.setFont(*author_style)
            
            c.setFillColor(colors.gray)
            c.drawString(left, y, authors_line)
            c.setFillColor(colors.black)
            y -= line_height
        
        y -= 4

    c.showPage()
    c.save()
    return pages_used


def build_volume(volume_num: int, total_volumes: int, entries: List[Entry], 
                 output_path: str, year: int, title: str, subtitle: str,
                 add_page_numbers: bool = False,
                 page_number_start: str = 'toc',
                 page_number_position: str = 'bottom-center') -> None:
    import tempfile
    
    def _compute_printed_starts(volume_entries: List[Entry], toc_pages: int) -> None:
        # Page numbering starts at 1 for first content page (after cover/blank/TOC)
        # TOC pages are unnumbered, so we track position for correct odd-page starts
        # but printed_start reflects content page numbers (1, 2, 3...)
        printed_counter = 0  # Content page counter
        for e in volume_entries:
            # Each paper should start on odd content page
            if printed_counter % 2 == 0 and printed_counter > 0:
                # Current position is even, need blank
                blanks = 1
            else:
                blanks = 0
            e.blanks_before = blanks
            e.printed_start = printed_counter + blanks + 1
            printed_counter += blanks + e.pages
    
    with tempfile.TemporaryDirectory() as tmpdir:
        cover_path = os.path.join(tmpdir, f'cover_v{volume_num}.pdf')
        toc_path = os.path.join(tmpdir, f'toc_v{volume_num}.pdf')
        
        cover_pages = generate_cover(cover_path, title, subtitle, year, volume_num, total_volumes)

        # Pass 1: assume 2 TOC pages (will be refined)
        _compute_printed_starts(entries, toc_pages=2)
        toc_pages = generate_toc(toc_path, entries, volume_num)
        # Pass 2: recompute with actual TOC pages and regenerate TOC
        _compute_printed_starts(entries, toc_pages=toc_pages)
        toc_pages = generate_toc(toc_path, entries, volume_num)
        
        writer = PdfWriter()
        
        # Add cover
        r = PdfReader(cover_path)
        for page in r.pages:
            writer.add_page(page)
        
        # Add blank page after cover
        writer.add_blank_page(width=A4[0], height=A4[1])
        blank_after_cover = 1
        
        # Add TOC
        r = PdfReader(toc_path)
        for page in r.pages:
            writer.add_page(page)
        
        # Ensure content starts on odd page
        total_before_content = cover_pages + blank_after_cover + toc_pages
        if total_before_content % 2 == 1:
            # Next page is even, add another blank
            writer.add_blank_page(width=A4[0], height=A4[1])
            total_before_content += 1
        
        # Add papers with optional blanks and bookmarks
        current_doc_pages = total_before_content
        for e in entries:
            # Insert blanks to ensure odd-page start per paper
            for _ in range(e.blanks_before):
                writer.add_blank_page(width=A4[0], height=A4[1])
                current_doc_pages += 1
            # Bookmark at this paper's first page (0-based index)
            writer.add_outline_item(e.title, page_number=current_doc_pages)
            # Append paper pages
            r = PdfReader(e.pdf_path)
            for page in r.pages:
                writer.add_page(page)
            current_doc_pages += len(r.pages)

        # Overlay duplex page numbers starting from first content page = 1; cover/blanks/TOC unnumbered
        if add_page_numbers:
            total_pages = len(writer.pages)
            for idx in range(total_before_content, total_pages):
                printed_no = (idx - total_before_content) + 1
                page = writer.pages[idx]
                w = float(page.mediabox.width)
                h = float(page.mediabox.height)

                buf = BytesIO()
                c = canvas.Canvas(buf, pagesize=(w, h))
                c.setFillColor(colors.black)
                c.setFont('Helvetica', 10)
                margin = 15 * mm
                text = str(printed_no)
                y = margin * 0.6
                if printed_no % 2 == 1:
                    # Odd page number -> bottom-right
                    c.drawRightString(w - margin, y, text)
                else:
                    # Even page number -> bottom-left
                    c.drawString(margin, y, text)
                c.showPage()
                c.save()
                overlay_page = PdfReader(BytesIO(buf.getvalue())).pages[0]
                page.merge_page(overlay_page)
        
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        with open(output_path, 'wb') as f:
            writer.write(f)


def _apply_page_numbers(writer: PdfWriter, *, cover_pages: int, toc_pages: int,
                        start_from: str = 'toc', position: str = 'bottom-center',
                        font_name: str = 'Helvetica', font_size: int = 10,
                        bottom_margin_pt: float = 14.0) -> None:
    """Overlay page numbers onto pages of writer.
    start_from: 'cover' | 'toc' | 'content'
    position: 'bottom-center' | 'bottom-left' | 'bottom-right'
    """
    total = len(writer.pages)
    if total == 0:
        return
    if start_from == 'cover':
        start_index = 0
    elif start_from == 'content':
        start_index = cover_pages + toc_pages
    else:  # 'toc'
        start_index = cover_pages

    for idx in range(start_index, total):
        display_no = idx - start_index + 1
        page = writer.pages[idx]
        mediabox = page.mediabox
        width = float(mediabox.width)
        height = float(mediabox.height)

        buf = BytesIO()
        c = canvas.Canvas(buf, pagesize=(width, height))
        c.setFillColor(colors.black)
        c.setFont(font_name, font_size)

        text = str(display_no)
        y = bottom_margin_pt
        if position == 'bottom-left':
            x = 20  # 20pt from left
            c.drawString(x, y, text)
        elif position == 'bottom-right':
            x = width - 20
            c.drawRightString(x, y, text)
        else:
            x = width / 2.0
            c.drawCentredString(x, y, text)
        c.showPage()
        c.save()

        overlay_reader = PdfReader(BytesIO(buf.getvalue()))
        overlay_page = overlay_reader.pages[0]
        # Merge overlay onto the existing page
        page.merge_page(overlay_page)


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description='Split large proceedings into multiple volumes')
    ap.add_argument('--csv', required=True, help='Path to metadata CSV')
    ap.add_argument('--output-dir', default='data/volumes', help='Output directory for volumes')
    ap.add_argument('--volumes', type=int, default=4, help='Target number of volumes (default: 4)')
    ap.add_argument('--title', default=None, help='Main title (default: AMOS <year> Proceedings)')
    ap.add_argument('--subtitle', default=None, help='Subtitle')
    ap.add_argument('--add-page-numbers', action='store_true', help='Add page numbers to each page (default: off)')
    ap.add_argument('--page-number-start', choices=['cover','toc','content'], default='toc', help='Where numbering starts (default: toc)')
    ap.add_argument('--page-number-position', choices=['bottom-center','bottom-left','bottom-right'], default='bottom-center', help='Page number position (default: bottom-center)')
    args = ap.parse_args(argv)

    year, entries = read_entries(args.csv)
    if not entries:
        print('[ERROR] No entries found')
        return 1
    
    measure_pages(entries)
    total_pages = sum(e.pages for e in entries)
    total_volumes = assign_volumes(entries, args.volumes)
    
    main_title = args.title or f'AMOS {year} Proceedings'
    sub_title = args.subtitle or f'AMOS Technical Papers – {year}'
    
    print(f"总论文数: {len(entries)} 篇")
    print(f"总页数: {total_pages} 页")
    print(f"分为 {total_volumes} 册 (目标 {args.volumes} 册)，每篇论文完整在一册内...\n")
    
    for vol in range(1, total_volumes + 1):
        vol_entries = [e for e in entries if e.volume == vol]
        vol_pages = sum(e.pages for e in vol_entries)
        output_path = os.path.join(args.output_dir, f'AMOS_Proceedings_{year}_Vol{vol}_of_{total_volumes}.pdf')
        
        print(f"  第 {vol} 册: {len(vol_entries)} 篇论文, {vol_pages} 页 -> {output_path}")
        build_volume(vol, total_volumes, vol_entries, output_path, year, main_title, sub_title,
                     add_page_numbers=args.add_page_numbers,
                     page_number_start=args.page_number_start,
                     page_number_position=args.page_number_position)
    
    print(f"\n✅ 完成！已生成 {total_volumes} 册，保存在 {args.output_dir}/")
    return 0


if __name__ == '__main__':
    raise SystemExit(main())
